#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, re, os
import time
from datetime import date
from sheepomatic_io import Csv_file,Sheep_brief_file,Shepherd_brief_file
from sheepomatic import Sheep,Shepherd
import sheepomatic

class Sheep_matcher:
   """Class defining sheep-shepherd matching"""
   @classmethod
   def match_sheep(cls,sheep_list,shepherd_list,hard_couples_hash,output_dir):
      """The main sheep-shepherd matching method
 
      Arguments:
      sheep_list -- list of sheep to match with shepherds
      shepherd_list -- list of shepherds to match with sheep
      hard_couples_hash -- list of already given sheep-shepherd pairs
      output_dir -- the directory to use for all output files
 
      The matching strategy can be changed based on the actual situation
      (e.g. more shepherds than sheep; more sheep than shepherds, date problems, etc.).
      The original strategy proceeds as follows:
      read hard-coded pairs
      first rounds (fixed number of times):
      match sheep faculty-by-faculty
      optionally: match odd sheep/shepherds from the faculty pools
      further rounds (until pairs still get added):
      match all sheep/shepherds
      fallback rounds (optionally, until pairs still get added):
      use a different (more relaxed) weighting function to add more pairs
      """
      #
      pair_odd_in_first_rounds = False
      use_fallback_rounds = True
      # INIT HASHES
      sheep_hash = {}
      shepherd_hash = {}
      #
      all_couples = []
      ava_sheep = []
      ava_shepherds = []
      #
      all_faculties = set([])
      #
      for sheep in sheep_list:
         sheep_hash[sheep.mail] = sheep
         ava_sheep.append(sheep.mail)
         all_faculties.add(sheep.faculty)
      for shepherd in shepherd_list:
         shepherd_hash[shepherd.mail] = shepherd
         ava_shepherds.append(shepherd.mail)
         all_faculties.add(shepherd.faculty)
         shepherd.aleady_has = 0
      if "unknown" in all_faculties:
         all_faculties.remove("unknown")
      #
      # HARD COUPLES
      for hard_sheep in hard_couples_hash:
         hard_shepherd = hard_couples_hash[hard_sheep]
         all_couples.append((hard_sheep,hard_shepherd))
         print >>sys.stderr,"Hard-coded couple:", hard_sheep,hard_shepherd
         ava_sheep.remove(hard_sheep)
         shepherd_hash[hard_shepherd].capacity -= 1
         if shepherd_hash[hard_shepherd].capacity == 0:
            ava_shepherds.remove(hard_shepherd)
      #
      # FIRST ROUNDS (faculty-by-faculty)
      for round in [1,2]:
         all_unused_sheep = []
         all_unused_shepherds = []
         # divide sheep and shepherds according to their faculties
         fac_sheep = {}
         for fac in all_faculties:
            fac_sheep[fac] = []
         fac_sheep["unknown"] = []
         for sheep in ava_sheep:
            fac_sheep[ sheep_hash[sheep].faculty ].append(sheep)
         #
         fac_shepherds = {}
         for fac in all_faculties:
            fac_shepherds[fac] = []
         for shepherd in ava_shepherds:
            fac_shepherds[ shepherd_hash[shepherd].faculty ].append(shepherd)
         #
         # pairing by faculties
         for faculty in all_faculties:
            couples,unused_sheep,unused_shepherds =\
                       sheepomatic.couple_sheep(fac_sheep[faculty],fac_shepherds[faculty],\
                                                sheep_hash,shepherd_hash,cls.pair_weight)
            all_couples.extend(couples)
            all_unused_sheep.extend(unused_sheep)
            all_unused_shepherds.extend(unused_shepherds)
            print >>sys.stderr, "round",round, faculty,\
                                "couples="+str(len(couples)),\
                                "sheep left="+str(len(unused_sheep)),\
                                "shepherds left="+str(len(unused_shepherds))
         # pairing the rest (unknown faculty sheep + all sheep/shepherds unused so far)
         all_unused_sheep.extend(fac_sheep["unknown"])
         if pair_odd_in_first_rounds:
            couples,unused_sheep,unused_shepherds =\
                          sheepomatic.couple_sheep(all_unused_sheep,all_unused_shepherds,\
                                                   sheep_hash,shepherd_hash,cls.pair_weight)
            all_couples.extend(couples)
            all_unused_sheep = unused_sheep
            all_unused_shepherds = unused_shepherds
            print >>sys.stderr, "round",round, "other",\
                                "couples="+str(len(couples)),\
                                "sheep left="+str(len(unused_sheep)),\
                                "shepherds left="+str(len(unused_shepherds))
         print >> sys.stderr, "all couples="+str(len(all_couples))
         print >> sys.stderr, "all sheep_left="+str(len(all_unused_sheep))
         print >> sys.stderr, "all shepherds_left="+str(len(all_unused_shepherds))
         # update ava_sheep and ava_shepherds
         ava_sheep = []
         for sheep in all_unused_sheep:
            ava_sheep.append(sheep)
         sheepomatic.update_ava_shepherds(ava_shepherds,all_unused_shepherds,shepherd_hash)
      # FURTHER ROUNDS (global) ... until new couples still get added ...
      last_found = 1
      round = 0
      while last_found > 0 and len(ava_sheep) > 0:
         round += 1
         couples,unused_sheep,unused_shepherds =\
                     sheepomatic.couple_sheep(ava_sheep,ava_shepherds,\
                                              sheep_hash,shepherd_hash,cls.pair_weight)
         all_couples.extend(couples)
         all_unused_sheep = unused_sheep
         all_unused_shepherds = unused_shepherds
         print >>sys.stderr, "all-"+str(round),\
                             "couples="+str(len(couples)),\
                             "sheep left="+str(len(unused_sheep)),\
                             "shepherds left="+str(len(unused_shepherds))
         # update ava_sheep and ava_shepherds
         ava_sheep = []
         for sheep in all_unused_sheep:
            ava_sheep.append(sheep)
         sheepomatic.update_ava_shepherds(ava_shepherds,all_unused_shepherds,shepherd_hash)
         # check fixed point
         last_found = len(couples)
      # FALLBACK ROUNDS (relaxing constraints)
      if use_fallback_rounds:
         last_found = 1
         round = 0
         while last_found > 0 and len(ava_sheep) > 0:
            round += 1
            couples,unused_sheep,unused_shepherds =\
                        sheepomatic.couple_sheep(ava_sheep,ava_shepherds,\
                                                 sheep_hash,shepherd_hash,cls.fallback_pair_weight)
            all_couples.extend(couples)
            all_unused_sheep = unused_sheep
            all_unused_shepherds = unused_shepherds
            print >>sys.stderr, "fallback-"+str(round),\
                                "couples="+str(len(couples)),\
                                "sheep left="+str(len(unused_sheep)),\
                                "shepherds left="+str(len(unused_shepherds))
            # update ava_sheep and ava_shepherds
            ava_sheep = []
            for sheep in all_unused_sheep:
               ava_sheep.append(sheep)
            sheepomatic.update_ava_shepherds(ava_shepherds,all_unused_shepherds,shepherd_hash)
            # check fixed point
            last_found = len(couples)
      # READY FOR OUTPUT
      # create directory and move old files (if necessary)
      if os.path.exists(output_dir):
         if os.path.isdir(output_dir):
            if os.path.exists(output_dir+"/"+"couples.txt"):
               os.rename(output_dir+"/"+"couples.txt",output_dir+"/"+"couples.old")
            if os.path.exists(output_dir+"/"+"unassigned.txt"):
               os.rename(output_dir+"/"+"unassigned.txt",output_dir+"/"+"unassigned.old")
         else:
            print >>sys.stderr, "Error writing result.", output_dir,"already exists as a file!"
            sys.exit(1)
      else:
         os.makedirs(output_dir)
      # write couples
      f1 = open(output_dir+"/"+"couples.txt", "w")
      for sheep,shepherd in all_couples:
         happy = 0
         if sheep_hash[sheep].country in set(shepherd_hash[shepherd].country_prefs):
            happy = 1
         f1.write(sheep + "," + shepherd + "," + \
                  str(happy) + "," + \
                  str(cls.pair_weight(sheep_hash[sheep],shepherd_hash[shepherd])) + "\n")
      f1.close()
      # write unassigned sheep
      f2 = open(output_dir+"/"+"unassigned_sheep.txt", "w")
      for sheep in ava_sheep:
         f2.write(sheep + "\n")
      f2.close()
      return

   @classmethod
   def pair_weight(cls,sheep,shepherd):
      """standard method useful when there are more sheep than shepherds
      and there is a date problem

      weight = number of common languages + (3 if country preference match)
      +1 if the arrival dates are similar
      -1 of the arrival dates differ a lot (an early tutors would be wasted on a late sheep)
      """
      # number of common languages
      clang = len(set(sheep.langs) & set(shepherd.langs))
      # country preferences
      cpref = 0
      if sheep.country in set(shepherd.country_prefs):
         cpref = 1
      # date constraints
      daydiff = (sheep.arrival_date - shepherd.ava_date).days
      if daydiff < 0 or clang == 0 or (not cls.is_sensitive_pair(sheep,shepherd)):
         return 0
      if daydiff < 2:
         return (3*cpref) + clang + 1
      if daydiff > 7:
         if ((3*cpref) + clang > 2):
            return (3*cpref) + clang - 1
         else:
            return 1
      return (3*cpref) + clang

   @classmethod
   def fallback_pair_weight(cls,sheep,shepherd):
      """standard fallback method when there is a date problem
      assigns 1 even to pairs where the tutor arrives later than the sheep
      (given the sheep is not travelling too far, especially not from overseas)
      """
      # number of common languages
      clang = len(set(sheep.langs) & set(shepherd.langs))
      # country preferences
      cpref = 0
      if sheep.country in set(shepherd.country_prefs):
         cpref = 1
      # date constraints
      daydiff = (sheep.arrival_date - shepherd.ava_date).days
      if daydiff < -8 or clang == 0 or (not cls.is_sensitive_pair(sheep,shepherd)):
         return 0
      if daydiff < 0 and sheep.country not in cls.close_countries():
         return 0
      return 1

   @classmethod
   def close_countries(cls):
      """Returns a list of countries close enough"""
      return set(["AT","BE","BG","HR","CZ","DK","DE","HU","IT","LV","LT","LU","NL",\
                  "PL","RO","RU","SK","SI","EE","FI","MK","ME","NO","RS","SE","UA"])

   @classmethod
   def is_sensitive_pair(cls,sheep,shepherd):
      """Checks whether the pair does not pose a cultural problem
      returns 0 on pairs from Mulslim countries and Japan
      where the sheep and spepherd gender wouldn't match
      """
      if sheep.country in set(["EG","SD","DZ","MA","IQ","IR","SA","YE","SY","TN",\
                               "SO","LY","JO","LB","PS","AE","MR","KW","OM","QA",\
                               "DJ","BH","KM","EH","TR","JP"]) and \
                   sheep.gender != "unknown" and \
                   sheep.gender != shepherd.gender:
         return 0
      return 1

def main(args):
    """Runs sheep-shepherd matching on the given files

    Parameters:
    sheep_file -- the sheep brief file to use
    shepherd_dile -- the shepherd brief file to use
    hard_couples_file -- the file with manually determined sheep-shepherd pairs
    output_dir -- the directory to write the output files to
    """
    # review parameters
    sheep_file = args[0]
    shepherd_file = args[1]
    hard_couples_file = args[2]
    output_dir = args[3]
    # read sheep + shepherds
    normed_sheep_list = Sheep_brief_file.read_sheep_brief_file(sheep_file)
    normed_shepherds_list = Shepherd_brief_file.read_shepherd_brief_file(shepherd_file)
    # read hard couples
    hard_couples_hash = sheepomatic.normfile2hash(hard_couples_file)
    # perform the matching (including writing output)
    Sheep_matcher.match_sheep(normed_sheep_list,normed_shepherds_list,hard_couples_hash,output_dir)

if __name__ == '__main__':
    args = sys.argv[1:]
    if len(args) == 4:
        main(args)
    else:
        print >> sys.stderr, "Wrong number of parameters."
        print >> sys.stderr, "Usage:"
        print >> sys.stderr, "sheep_matcher.py sheep_brief_file shepherd_brief_file hard_couples_file output_dir"
        sys.exit(1)
